#!/usr/bin/python3
"""
Partition a target using parted(8)

Use the `parted`(8) command line utility to partition the specified
`device`.

Buildhost commands used: `parted`.
"""

import subprocess
import sys


import osbuild.api


SCHEMA_2 = r"""
"devices": {
  "type": "object",
  "additionalProperties": true,
  "required": ["device"],
  "properties": {
    "device": {
      "type": "object",
      "additionalProperties": true
    }
  }
},
"options": {
  "additionalProperties": false,
  "required": ["label"],
  "properties": {
    "label": {
      "description": "The type of the partition table",
      "type": "string",
      "enum": ["gpt"]
    },
    "partitions": {
      "description": "Partition layout ",
      "type": "array",
      "required": ["start", "size", "name"],
      "items": {
        "description": "Description of one partition",
        "type": "object",
        "properties": {
          "bootable": {
            "description": "Mark the partition as bootable",
            "type": "boolean"
          },
          "name": {
            "description": "The partition name (GPT)",
            "type": "string"
          },
          "size": {
            "description": "The size of this partition",
            "type": "integer"
          },
          "start": {
            "description": "The start offset of this partition",
            "type": "integer"
          },
          "type": {
            "description": "The partition type",
            "type": "string"
          }
        }
      }
    }
  }
}
"""


class Partition:
    def __init__(self,
                 pttype: str = None,
                 start: int = None,
                 size: int = None,
                 bootable: bool = False,
                 name: str = None,
                 uuid: str = None):
        self.type = pttype
        self.start = start
        self.size = size
        self.bootable = bootable
        self.name = name
        self.uuid = uuid
        self.index = None

    @property
    def start_in_bytes(self):
        return (self.start or 0) * 512

    @property
    def size_in_bytes(self):
        return (self.size or 0) * 512

    @classmethod
    def from_json(cls, js) -> "Partition":
        p = cls(pttype=js.get("type"),
                start=js.get("start"),
                size=js.get("size"),
                bootable=js.get("bootable"),
                name=js.get("name"),
                uuid=js.get("uuid"))
        return p


class PartitionTable:
    def __init__(self, label, uuid, partitions):
        self.label = label
        self.uuid = uuid
        self.partitions = partitions or []

    def __getitem__(self, key) -> Partition:
        return self.partitions[key]

    def write_to(self, target):
        """Write the partition table to disk"""
        # generate the command for parted to create the table
        commands = [
            "unit", "s",
            "mklabel", self.label
        ]

        for i, p in enumerate(self.partitions):

            commands += [
                "mkpart", f'"{p.name}"', str(p.start), str(p.start + p.size - 1)
            ]

            if p.bootable:
                commands += ["set", str(i+1), "boot", "on"]

            if p.type:
                commands += ["set", str(i+1), p.type, "on"]

        subprocess.run(["parted", "-a", "none", "-s",
                        target, "--"] + commands,
                       encoding='utf-8',
                       check=True)

        subprocess.run(["parted", "-a", "none", "-s",
                        target, "--", "print"],
                       encoding='utf-8',
                       check=True)


def main(devices, options):
    device = devices["device"]["path"]

    pttype = options["label"]
    partitions = options.get("partitions")

    parts = [Partition.from_json(p) for p in partitions]
    pt = PartitionTable(pttype, None, parts)

    pt.write_to(device)


if __name__ == '__main__':
    args = osbuild.api.arguments()
    ret = main(args["devices"], args["options"])
    sys.exit(ret)
